import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

def get_fc_discriminator(num_classes, ndf=64):
    return discriminator()

def get_fc_discriminator2(num_classes, ndf=64):
    return discriminator2()


class discriminator(nn.Module):

    def __init__(self, ndf=64):
        super().__init__()

        self.conv1 = nn.Sequential(
        nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=2, padding=1),
        #nn.AdaptiveAvgPool2d((1, 1)),
    )


    def forward(self, x):
        output = self.conv1(x)
        #output = output.view(output.size(0), -1)
        #output = self.fc(output)
        #print(output.shape)
        return output


class discriminator2(nn.Module):

    def __init__(self, ndf=1):
        super().__init__()

        self.conv1 = nn.Sequential(
        nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=2, padding=1),
    )


    def forward(self, x):
        output = self.conv1(x)
        return output


class discriminator3(nn.Module):

    def __init__(self, ndf=64):
        super().__init__()

        self.conv1 = nn.Sequential(
        nn.Conv2d(2, ndf, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.AdaptiveAvgPool2d((1, 1)),
    )

        self.fc = nn.Linear(ndf * 8, 1)


    def forward(self, x):
        output = self.conv1(x)
        output = output.view(output.size(0), -1)
        output = self.fc(output)
        return output